from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import uvicorn
from typing import Dict, List, Set
import json
import time

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        # 使用字典存储不同聊天室的连接
        self.room_connections: Dict[str, Set[WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, room_id: str):
        await websocket.accept()
        # 如果聊天室不存在则创建
        if room_id not in self.room_connections:
            self.room_connections[room_id] = set()
        self.room_connections[room_id].add(websocket)
        print(f"用户加入聊天室 {room_id}, 当前人数: {len(self.room_connections[room_id])}")
    
    def disconnect(self, websocket: WebSocket, room_id: str):
        if room_id in self.room_connections:
            self.room_connections[room_id].discard(websocket)
            print(f"用户离开聊天室 {room_id}, 剩余人数: {len(self.room_connections[room_id])}")
            # 如果聊天室空了，可以清理（可选）
            if len(self.room_connections[room_id]) == 0:
                del self.room_connections[room_id]
    
    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)
    
    async def broadcast(self, message: str, room_id: str):
        if room_id in self.room_connections:
            # 使用集合的副本遍历，避免在迭代时修改集合
            for connection in self.room_connections[room_id].copy():
                try:
                    await connection.send_text(message)
                except Exception as e:
                    print(f"发送消息时出错: {e}")
                    self.disconnect(connection, room_id)

manager = ConnectionManager()

@app.get("/")
async def get():
    return HTMLResponse("""
    <html>
        <head><title>多聊天室WebSocket服务器</title></head>
        <body>
            <h1>多聊天室WebSocket服务器</h1>
            <p>请使用前端页面连接，支持多个聊天室</p>
        </body>
    </html>
    """)

# 修改WebSocket端点，支持room_id参数
@app.websocket("/ws/{room_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: str):
    await manager.connect(websocket, room_id)
    
    try:
        while True:
            data = await websocket.receive_text()
            
            try:
                json_data = json.loads(data)
                message = {
                    "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                    "room": room_id,
                    "client": f"用户{len(manager.room_connections.get(room_id, set()))}",
                    "content": json_data.get("content", ""),
                    "type": "chat_message"
                }
                
                await manager.broadcast(json.dumps(message, ensure_ascii=False), room_id)
                
            except json.JSONDecodeError:
                error_msg = {
                    "status": "error",
                    "message": "必须发送JSON格式数据",
                    "timestamp": int(time.time())
                }
                await websocket.send_text(json.dumps(error_msg))
                
    except WebSocketDisconnect:
        manager.disconnect(websocket, room_id)
        message = {
            "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
            "room": room_id,
            "client": "系统",
            "content": "有用户离开了聊天室",
            "type": "system_message"
        }
        await manager.broadcast(json.dumps(message, ensure_ascii=False), room_id)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)